import engine
import random
import pygame
import json
import math

import pymunk
import pymunk.constraints
import pymunk.pygame_util
import numpy as np


class Enviroment:
    def __init__(self, map_path:str):
        pygame.init()
        self.env = None
        self.screen = pygame.Surface((700, 600))
        self.display = None
        self.draw_options = pymunk.pygame_util.DrawOptions(self.screen)
        self.wall_poslist = None
        self.insert_limit = 50
        self.max_steps = 300	
        self.now_steps = 0
        self.bg = None
        self.radius = 7
        with open(map_path) as json_file:
            self.wall_poslist = json.load(json_file)	

    def step(self, action):
        self.now_steps += 1
        exact = 1000
        action = int(action)
        # 输入动作
        if action == 0:
            if len(self.env.balls) <= self.insert_limit:
                self.env.insert_a_ball(radius=self.radius)
        elif action == 1:
            self.env.pull_back_a_ball()
        elif action == 2:
            self.env.bend_guidewire(d_angle=10)
        elif action == 3:
            self.env.bend_guidewire(d_angle=-10)
        elif action == 4:
            pass
        else:
            print(f"动作“{action}”不合法")
            raise "动作不合法"
        
        # 物理引擎进行计算
        for x in range(exact):
            self.env.space.step(1 / exact)	
        
        # 导丝以及终点的位置坐标
        debug = self.env.balls[-1].body.position.int_tuple	
        debug2 = self.env.goals[-1].body.position.int_tuple
        
        reward = 0
        done = False
        if self.env.done:
            done = True	
        self.screen.fill(pygame.Color("white"))          # 清空图案
        self.env.space.debug_draw(self.draw_options)  # 绘制新图案
        
        w, h = 40, 40
        rect = (debug2[0] - w/2, debug2[1] - h/2, w, h)
        pygame.draw.rect(self.screen, (0, 0, 0), rect, 5)
        rect=(debug[0] - w/2, debug[1] - h/2, w, h)
        pygame.draw.rect(self.screen, (255, 0, 0), rect, 5)
        s = np.array(pygame.surfarray.pixels3d(self.screen))
        
        if self.now_steps >= self.max_steps:
            done = True
        if done:
            distance = pow((pow(debug[0] - debug2[0], 2) + pow(debug[1] - debug2[1], 2)), 0.5) * 0.001		
            reward = - math.log(distance)
        return s, reward, done, (debug, debug2, self.now_steps)

    def reset(self, g_pos=(270,100), keep=False):
        if not keep:
        # 清除环境内的所有东西
            del self.env
            self.env = engine.guide_wire_space()
            # 载入血管壁
            self.env.draw_walls(self.wall_poslist)
            # 载入导丝
            self.env.create_goal(g_pos)
            self.env.create_ini_guidewire(radius=self.radius,ini_len=5)
        else:
            old_balls = self.env.balls
            del self.env
            self.env = engine.guide_wire_space()
            # 载入血管壁
            self.env.draw_walls(self.wall_poslist)
            # 载入导丝
            self.env.create_goal(g_pos)
            self.env.draw_old_guidewire(old_balls, radius=self.radius)
        s, *_ = self.step(4)
        self.now_steps = 0
        return s

    def render(self):
        """唤起窗口渲染画面, 返回窗口是否被用户关闭的句柄"""
        if self.display is None:
            self.display = pygame.display.set_mode((700, 600))
        
        self.display.blit(self.screen, (0,0))
        pygame.display.flip()
        pygame.display.set_caption(f"step:{self.now_steps}/{self.max_steps}")
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return True
        return False

    def beautiful_render(self):
        """唤起窗口渲染画面, 返回窗口是否被用户关闭的句柄"""
        if self.display is None:
            self.display = pygame.display.set_mode((700, 600))
        
        self.display.blit(self.bg, (0,0))
        points = []
        
        # 获取导丝的位姿信息并绘制导丝
        for i in self.env.balls:
            points.append(i.body.position.int_tuple)
        pygame.draw.lines(self.display, (110, 110, 100), False, points, width=5)
        pygame.draw.circle(self.display, (150, 150, 150), points[-1], 3)
        
        # 绘制肿瘤
        for i in self.env.goals:		
            p = self.env.goals[-1].body.position.int_tuple
            pygame.draw.circle(self.display, (70, 27, 44), p, 15)
        
        # 绘制人工智能识别的结果
        w, h = 40, 40
        rect = (points[-1][0] - w/2, points[-1][1] - h/2, w, h)
        pygame.draw.rect(self.display, (0, 255, 0), rect, 5)
        rect = (p[0] - w/2, p[1] - h/2, w, h)
        pygame.draw.rect(self.display, (255, 255, 255), rect, 5)


        pygame.display.flip()
        pygame.display.set_caption(f"step:{self.now_steps}/{self.max_steps}")
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return True
        return False


    def close(self):
        if self.display is not None:
            pygame.display.quit()
            self.display = None

def main():
	gym = Enviroment()
	# 读取血管路径
	with open("600700.json") as json_file:
		gym.wall_poslist = json.load(json_file)	
	gym.bg = pygame.image.load("img_rsc\\3.png")
	
	# 初始化环境
	s = gym.reset((200, 350))
	print(type(s))
	total = 0
	for i in range(500):
		s, r, d, debug = gym.step(random.randint(0,4))
		print(i, r, d, debug)
		gym.render()
		total += r
		if d:
			break
	print(total)
	gym.close()
               
	s = gym.reset()
	total = 0
	for i in range(500):
		s, r, d, debug = gym.step(random.randint(0,4))
		print(i, r, d, debug)
		gym.beautiful_render()
		total += r
		if d:
			break
	gym.close()
	print(total)



if __name__ == '__main__':
	main()
 
